package cn.com.zetatech.loader.common.db;

import cn.com.zetatech.loader.api.exception.LoaderException;
import cn.hutool.json.JSONUtil;
import lombok.Data;
import lombok.ToString;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.codec.binary.Base64;
import org.apache.commons.lang3.StringUtils;
import org.apache.http.HttpHeaders;
import org.apache.http.client.methods.CloseableHttpResponse;
import org.apache.http.client.methods.HttpPut;
import org.apache.http.entity.StringEntity;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.impl.client.DefaultRedirectStrategy;
import org.apache.http.impl.client.HttpClientBuilder;
import org.apache.http.impl.client.HttpClients;
import org.apache.http.util.EntityUtils;

import java.nio.charset.StandardCharsets;
import java.util.UUID;

/**
 * StarRocks通过Stream CloseableHttpClient存储数据
 */
@Slf4j
public class StarRocksStreamLoad {

    /**
     * @param starrocksConfigInfo starrocks配置信息
     * @param tableName           待入库表名
     * @param columns             待入库表名对应的字段名（若为空，则数据必须按表中字段顺序组装）
     * @param content             数据以csv文件类型组装，分隔符为\t。换行符为\n
     * @throws Exception
     */
    public void sendData(StarrocksConfigInfo starrocksConfigInfo, String tableName, String columns, String content)
            throws Exception {
        if (starrocksConfigInfo == null) {
            throw new LoaderException("the datasource is invalid dbConfig!");
        }

        String STARROCKS_HOST = starrocksConfigInfo.getIp();
        int STARROCKS_HTTP_PORT = starrocksConfigInfo.getPort();
        String STARROCKS_DB = starrocksConfigInfo.getDbScheme();
        String STARROCKS_TABLE = tableName;
        String STARROCKS_USER = starrocksConfigInfo.getUsername();
        String STARROCKS_PASSWORD = starrocksConfigInfo.getPassword();

        // Stream Load用于向指定的表导入数据，它可以保证每次导入任务的原子性（整批数据要么全部导入成功，要么全部失败），其完整语法为：
        // curl --location-trusted -u user:passwd [-H ""...] -T data.file -XPUT http://fe_host:http_port/api/{db}/{table}/_stream_load
        final String loadUrl = String.format("http://%s:%s/api/%s/%s/_stream_load",
                STARROCKS_HOST,
                STARROCKS_HTTP_PORT,
                STARROCKS_DB,
                STARROCKS_TABLE);

        final HttpClientBuilder httpClientBuilder = HttpClients
                .custom()
                .setRedirectStrategy(new DefaultRedirectStrategy() {
                    @Override
                    protected boolean isRedirectable(String method) {
                        return true;
                    }
                });
        try (CloseableHttpClient client = httpClientBuilder.build()) {
            HttpPut put = new HttpPut(loadUrl);
            StringEntity entity = new StringEntity(content, "UTF-8");
            put.setHeader(HttpHeaders.EXPECT, "100-continue");
            put.setHeader(HttpHeaders.AUTHORIZATION, basicAuthHeader(STARROCKS_USER, STARROCKS_PASSWORD));
//            put.setHeader("column_separator", "\t"); // 默认分割符:\t
//            put.setHeader("row_delimiter", "\n"); // 默认分割符:\n
            if (!StringUtils.isEmpty(columns)) {
                put.setHeader("columns", columns);
            }
            // the label header is optional, not necessary. use label header can ensure at most once semantics
            put.setHeader("label", createBatchLabel(null));
            put.setEntity(entity);

            try (CloseableHttpResponse response = client.execute(put)) {
                String loadResult = "";
                if (response.getEntity() != null) {
                    loadResult = EntityUtils.toString(response.getEntity());
                }
                final int statusCode = response.getStatusLine().getStatusCode();
                // statusCode 200 just indicates that starrocks be service is ok, not stream load
                // you should see the output content to find whether stream load is success
                if (statusCode != 200) {
                    throw new LoaderException(
                            String.format("Stream load failed, statusCode=%s load result=%s", statusCode, loadResult));
                }
                StreamLoadResult streamLoadResult = JSONUtil.toBean(loadResult, StreamLoadResult.class);
                log.info("Stream load result: {}", streamLoadResult.toString());
                if (!"Success".equalsIgnoreCase(streamLoadResult.getStatus())) {
                    throw new LoaderException(String.format("Stream load failed, status=%s load result=%s", streamLoadResult.getStatus(), streamLoadResult.getMessage()));
                }
            } catch (Exception e) {
                throw new Exception(e.getMessage());
            }
        }
    }

    public static String basicAuthHeader(String username, String password) {
        final String tobeEncode = username + ":" + password;
        byte[] encoded = Base64.encodeBase64(tobeEncode.getBytes(StandardCharsets.UTF_8));
        return "Basic " + new String(encoded);
    }

    public static String createBatchLabel(String labelPrefix) {
        StringBuilder sb = new StringBuilder();
        if (!StringUtils.isEmpty(labelPrefix)) {
            sb.append(labelPrefix);
        }
        return sb.append(UUID.randomUUID().toString()).toString();
    }

    public static void main(String[] args) throws Exception {
        int id1 = 1;
        int id2 = 10;
        String id3 = "Simon";
        int rowNumber = 10;
        String oneRow = id1 + "\t" + id2 + "\t" + id3 + "\n";

        StringBuilder stringBuilder = new StringBuilder();
        for (int i = 0; i < rowNumber; i++) {
            stringBuilder.append(oneRow);
        }

        stringBuilder.deleteCharAt(stringBuilder.length() - 1);

        String loadData = stringBuilder.toString();
        StarRocksStreamLoad starrocksStreamLoad = new StarRocksStreamLoad();
//        starrocksStreamLoad.sendData(loadData);

        System.out.println(basicAuthHeader("root", ""));
    }

    @ToString
    @Data
    public static class StreamLoadResult {

        /**
         * 导入作业的事务 ID。
         */
        private int TxnId;

        /**
         * 导入作业的标签。
         */
        private String Label;

        /**
         * 此次导入的数据的最终状态。
         * 可能的值: Success, Publish Timeout, Label Already Exists, Fail
         */
        private String Status;

        /**
         * 导入作业的状态详情。如果导入作业失败，这里会返回具体的失败原因。
         */
        private String Message;

        /**
         * 读取到的总行数。
         */
        private int NumberTotalRows;

        /**
         * 成功导入的总行数。只有当返回结果中的 Status 为 Success 时有效。
         */
        private int NumberLoadedRows;

        /**
         * 导入过程中因数据质量不合格而过滤掉的行数。
         */
        private int NumberFilteredRows;

        /**
         * 导入过程中根据 WHERE 子句指定的条件而过滤掉的行数。
         */
        private int NumberUnselectedRows;

        /**
         * 此次导入的数据量大小。单位：字节 (Bytes)。
         */
        private int LoadBytes;

        /**
         * 此次导入所用的时间。单位：毫秒 (ms)。
         */
        private int LoadTimeMs;

        /**
         * 导入作业开启事务的时长。单位：毫秒 (ms)。
         */
        private int BeginTxnTimeMs;

        /**
         * 导入作业生成执行计划的时长。单位：毫秒 (ms)。
         */
        private int StreamLoadPutTimeMS;

        /**
         * 导入作业读取数据的时长。单位：毫秒 (ms)。
         */
        private int ReadDataTimeMs;

        /**
         * 导入作业写入数据的时长。单位：毫秒 (ms)。
         */
        private int WriteDataTimeMs;

        /**
         * 导入作业提交和数据发布的耗时。单位：毫秒 (ms)。
         */
        private int CommitAndPublishTimeMs;
    }
}
